#include "GlassMaterial.cuh"

using namespace OBR;

static __forceinline__ __device__ float3 refract(const float3 &wi, const float3 &n, float eta)
{
    // Compute $\cos \theta_\roman{t}$ using Snell's law
    float cosThetaI = -dot(n, wi);
    float sin2ThetaI = fmax(float(0), float(1 - cosThetaI * cosThetaI));
    float sin2ThetaT = eta * eta * sin2ThetaI;

    // Handle total internal reflection for transmission
    if (sin2ThetaT >= 1)
    {
        return reflect(wi, n);
    }
    else
    {
        float cosThetaT = sqrtf(1 - sin2ThetaT);
        return normalize(eta * wi - (eta * cosThetaI - cosThetaT) * n);
    }
}

/**
 * @brief generate new ray info and add to the next ray, then return beta of the ray
 *
 * @param prd
 * @param instance_id
 * @return beta
 */
static __forceinline__ __device__ Spectrum rayGeneration(SamplePRD *prd, float *pdf,
                                                         GlassMaterialData *data)
{
    float eta = (dot(prd->inter.n, prd->inter.wo.d) < 0) ? data->eta : 1.0f / data->eta;
    float3 n = (dot(prd->inter.n, prd->inter.wo.d) < 0) ? prd->inter.n : -prd->inter.n;
    float reflectance = fr_dielectric(clamp(dot(-prd->inter.wo.d, n), 0.0f, 1.0f), eta);

#ifdef DEBUG_SHADER
    if (prd->log)
    {
        printf("reflectance: %f\n", reflectance);
    }
#endif

    if (curand_uniform(prd->rng) < reflectance)
    {
        prd->inter.wi = prd->inter.spawnRay(reflect(prd->inter.wo.d, n));
    }
    else
    {
        prd->inter.wi = prd->inter.spawnRay(refract(prd->inter.wo.d, n, eta));
    }

#ifdef DEBUG_SHADER
    if (prd->log)
    {
        printf("in: ");
        print_float3(prd->inter.wi.d);
        printf("\n");

        printf("normal: ");
        print_float3(n);
        printf("\n");

        printf("out: ");
        print_float3(prd->inter.wo.d);
        printf("\n");
    }
#endif

    *pdf = 1.0f;
    return Spectrum(1.0f);
}

extern "C" __device__ Spectrum __direct_callable__glass_material(SamplePRD *prd, float *pdf,
                                                                 float u0, float u1,
                                                                 MaterialCallMode mode)
{
    // get sbt data
    CallableData *sbt_record = reinterpret_cast<CallableData *>(optixGetSbtDataPointer());
    GlassMaterialData *sbt_data =
        reinterpret_cast<GlassMaterialData *>(sbt_record->callable_payload);
    sbt_data = sbt_data + prd->inter.mat_instance_id;

    if (mode == MaterialCallMode::SAMPLE_FR)
    {
        return rayGeneration(prd, pdf, sbt_data);
    }
    else if (mode == MaterialCallMode::CALC_FR_PDF)
    {
        *pdf = 0;
        return Spectrum(0.0f);
    }
    return Spectrum(0.0f);
}